package org.wenshu.ai.modular.chat.provider.core;

import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.jina.JinaEmbeddingModel;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.model.openai.OpenAiStreamingChatModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import jakarta.annotation.PostConstruct;
import lombok.Getter;
import org.springframework.stereotype.Service;
import org.wenshu.ai.config.AiModelConfig;
import org.wenshu.ai.modular.chat.provider.core.oceanbase.OceanbaseEmbeddingStore;

import static java.time.Duration.ofSeconds;

@Service
@Getter
public class AiModelService {

    private final AiModelConfig aiModelConfig;

    private EmbeddingStore<TextSegment> ddlEmbeddingStore;
    private EmbeddingStore<TextSegment> sqlEmbeddingStore;
    private EmbeddingStore<TextSegment> docEmbeddingStore;
    private OpenAiStreamingChatModel streamingChatModel;
    private ChatLanguageModel chatModel;
    private EmbeddingModel embeddingModel;

    public AiModelService(AiModelConfig aiModelConfig) {
        this.aiModelConfig = aiModelConfig;
    }

    @PostConstruct
    public void init() {
        // 初始化向量存储
        ddlEmbeddingStore = new OceanbaseEmbeddingStore(
            aiModelConfig.getVectorStore().getJdbcUrl(),
            aiModelConfig.getVectorStore().getUsername(),
            aiModelConfig.getVectorStore().getPassword(),
            aiModelConfig.getVectorStore().getCollections().getDdl(),
            aiModelConfig.getVectorStore().getDimensions()
        );

        sqlEmbeddingStore = new OceanbaseEmbeddingStore(
            aiModelConfig.getVectorStore().getJdbcUrl(),
            aiModelConfig.getVectorStore().getUsername(),
            aiModelConfig.getVectorStore().getPassword(),
            aiModelConfig.getVectorStore().getCollections().getSql(),
            aiModelConfig.getVectorStore().getDimensions()
        );

        docEmbeddingStore = new OceanbaseEmbeddingStore(
            aiModelConfig.getVectorStore().getJdbcUrl(),
            aiModelConfig.getVectorStore().getUsername(),
            aiModelConfig.getVectorStore().getPassword(),
            aiModelConfig.getVectorStore().getCollections().getDoc(),
            aiModelConfig.getVectorStore().getDimensions()
        );

        // 初始化对话模型
        streamingChatModel = OpenAiStreamingChatModel.builder()
            .baseUrl(aiModelConfig.getLlm().getBaseUrl())
            .apiKey(aiModelConfig.getLlm().getApiKey())
            .modelName(aiModelConfig.getLlm().getModelName())
            .temperature(aiModelConfig.getLlm().getTemperature())
            .timeout(ofSeconds(aiModelConfig.getLlm().getTimeoutSeconds()))
            .logRequests(aiModelConfig.getLlm().getLogRequests())
            .logResponses(aiModelConfig.getLlm().getLogResponses())
            .build();

        chatModel = OpenAiChatModel.builder()
            .baseUrl(aiModelConfig.getLlm().getBaseUrl())
            .apiKey(aiModelConfig.getLlm().getApiKey())
            .modelName(aiModelConfig.getLlm().getModelName())
            .temperature(aiModelConfig.getLlm().getTemperature())
            .timeout(ofSeconds(aiModelConfig.getLlm().getTimeoutSeconds()))
            .logRequests(aiModelConfig.getLlm().getLogRequests())
            .logResponses(aiModelConfig.getLlm().getLogResponses())
            .build();

        // 初始化向量模型
        embeddingModel = JinaEmbeddingModel.builder()
            .apiKey(aiModelConfig.getEmbedding().getApiKey())
            .modelName(aiModelConfig.getEmbedding().getModelName())
            .timeout(ofSeconds(aiModelConfig.getEmbedding().getTimeoutSeconds()))
            .lateChunking(aiModelConfig.getEmbedding().getLateChunking())
            .logRequests(aiModelConfig.getEmbedding().getLogRequests())
            .logResponses(aiModelConfig.getEmbedding().getLogResponses())
            .build();
    }
}